import pandas as pd
from transformers import TextClassificationPipeline, BertForSequenceClassification, BertTokenizer

# 使用训练好的fine-tune模型
from dataProcess import test_text, test_label
from utils.truncate import greater512

# 暂且使用test集合来做inference

model = BertForSequenceClassification.from_pretrained('./fine-tune-event')

# 直接把原来的bert模型中的tokenizer.json,token的config文件，vocab.txt
# 实际上这些就是从字，到cls+XXX+sep，到查表变成input_ids
tokenizer = BertTokenizer.from_pretrained('./bert-base-chinese')

pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)

# test = [
#     '方直科技(300235)总经理辞职一铭软件因2017年报仍未披露被提示摘牌风险',
#     '国民信托产品违约僵局待解 董事长建议投资者起诉公司海航质疑*ST九龙(600555)案判决 李勤夫方董事会或重新控制公司',
#     '成都中禾投资公司老板跑路 专家忧年底或现倒闭潮昆百大A(000560)重要股东累计减持5850万股',
#     '因大股东发生未披露事项 云内动力(000903)午后停牌惠发股份股东北京弘富拟减持不超3%股份'
# ]
id2query = {0: '其他', 1: '信批违规', 2: '实控人股东变更', 3: '交易违规', 4: '涉嫌非法集资', 5: '不能履职', 6: '重组失败', 7: '评级调整', 8: '业绩下滑',
            9: '涉嫌违法', 10: '财务造假', 11: '涉嫌传销', 12: '涉嫌欺诈', 13: '资金账户风险', 14: '高管负面', 15: '资产负面', 16: '投诉维权', 17: '产品违规',
            18: '提现困难', 19: '失联跑路', 20: '歇业停业', 21: '公司股市异常'}

review_list = []
predict_list = []
gold_list = []

for i, t in enumerate(test_text):
    t = greater512(t, 'ht')
    predict = pipeline(t)[0]['label']
    predict = id2query[int(predict[predict.index('_') + 1:])]
    gold = id2query[test_label[i]]
    if predict != gold:
        review_list.append(t)
        predict_list.append(predict)
        gold_list.append(gold)

res = pd.DataFrame()
res['review'] = pd.Series(review_list)
res['predict'] = pd.Series(predict_list)
res['gold'] = pd.Series(gold_list)

res.to_csv('wrong_analysis.csv')

# # inference的时候怎么截断呢？准备数据集的时候都截断了，在token每个数据集合的时候进行的操作，那么投入到实际使用当中应该怎么弄呢？
# 参考论文How to Fine-Tune BERT for Text Classification?
# 1. head-only: keep the first 510 tokens6;
# 2. tail-only: keep the last 510 tokens;
# 3. head+tail: empirically select the first 128 and the last 382 tokens.
# 三种方法，开门见山，保留开头的中心句；最后收尾总结，保留最后的总结句；兼顾开头+结尾
# for t in test_text:
#     if len(t) < 512:
#         print(t, pipeline(t)[0]['label'])
